import logging
import json
from Config import *
from pathlib import Path

from LoggerManager import LoggerManager
from CKGRetriever import CKGRetriever
from EnvironmentService import EnvironmentService
from Planner import PlannerState
from Planner import graph
from Utils import add_record, add_test_case

config = {"recursion_limit": RECURSION_LIMIT}
graph_retriever = CKGRetriever("bolt://localhost:7687", "neo4j", Neo4j_Password)
# envServer = EnvironmentService()
# envServer.set_base_path(r"E:\Project\Agent\TestAgent\TestAgent\TestAgent\CKGConstruction\test4")
methods_list = graph_retriever.load_methods()
# methods_list = methods_list[:2]
# if commons-lang and jfreechart, use this
# methods_list = graph_retriever.load_filtered_methods()

############### log setting ###############
log_dir = "result/log"
logger_manager = LoggerManager()
logger = logger_manager.logger
###########################################

if not os.path.exists(log_dir):
    os.makedirs(log_dir)
generate_record_path = os.path.join(log_dir, "record.jsonl")


def read_jsonl(file_path):
    data = []
    with open(file_path, "r") as file:
        for line in file:
            data.append(json.loads(line))
    return data

def run_generation(project_path):
    index = 0
    max_retries = 2
    for method in methods_list:
        retries = 0
        index += 1
        if not os.path.exists(f"{log_dir}/detail/{str(index) + '_' + method.name}"):
            os.makedirs(f"{log_dir}/detail/{str(index) + '_' + method.name}")

        log_path = f"{log_dir}/detail/{str(index) + '_' + method.name}/{index}_{method.name}.log"
        logger_manager.change_log_path(log_path)
        logger.info(str(index) + '_' + method.name + " " + method.signature)
        while retries < max_retries:
            try:
                logger.info("===================== TEST GENERATE START =====================\n")
                logger.info(f"package_name: {method.package_name}\n")
                logger.info(f"method_name: {method.name}\n")
                logger.info(f"method_signature: {method.signature}\n")
                logger.info(f"full_qualified_name: {method.full_qualified_name}\n")
                logger.info(f"absolute_path: {method.absolute_path}\n\n")
                graph_retriever.change_focal_method_id(int(method.id))

                method_path = Path(method.absolute_path)
                parts = list(method_path.parts)
                new_parts = [p if p != "main" else "test" for p in parts]
                inject_dir = str(Path(*new_parts).parent)
                envServer = EnvironmentService()
                envServer.set_base_path(project_path)
                envServer.set_inject_dir(inject_dir)
                clazz = graph_retriever.search_clazz_query(method.class_name,
                                                           method.full_qualified_name.rsplit('.', 1)[0])

                state = PlannerState(envServer=envServer,
                                     package_name=method.package_name,
                                     method_id=int(method.id),
                                     method_code=method.content,
                                     method_signature=method.signature,
                                     class_name=method.class_name,
                                     full_method_name=method.full_qualified_name,
                                     start_line=method.start_line,
                                     end_line=method.end_line)

                for event in graph.stream(state, config, subgraphs=True):
                    local_log_path = f"{log_dir}/detail/{str(index) + '_' + method.name}/{str(event[0])}.log"
                    local_handler = logging.FileHandler(local_log_path, mode='a')
                    logger_manager.add_file_handler(local_handler)

                    event = event[1]
                    EndPoint = 'testCaseAcceptor'
                    if EndPoint in event.keys():
                        logger.info("===================== TEST GENERATE END =====================\n\n")
                        result = []
                        test_code = []
                        find_bugs = []
                        covered_lines = set()
                        total_lines = set()
                        mutation_info = {}
                        line_coverage_rates = []
                        branch_coverage_rates = []
                        mutation_scores = []
                        mutants_list = []
                        test_point = []
                        test_report = []
                        inner_index = 0
                        for test_case in event[EndPoint]['test_cases']:
                            inner_index += 1
                            logger.info(f"===================== Test Report {inner_index} =====================\n")
                            logger.info(f"test case: {LoggerManager.format_dict(test_case['test_case'])}\n")
                            logger.info(f"test result: {LoggerManager.format_dict(test_case['test_result'])}\n")
                            logger.info(f"find bugs: {LoggerManager.format_dict(test_case['find_bug'])}\n")
                            logger.info(f"test point: {LoggerManager.format_dict(test_case['test_point'])}\n")
                            logger.info(f"test report: {LoggerManager.format_dict(test_case['test_report'])}\n")
                            logger.info(f"coverage report: {LoggerManager.format_dict(test_case['coverage_report'])}\n")
                            logger.info(f"mutation report: {LoggerManager.format_dict(test_case['mutation_report'])}\n")
                            result.append(test_case['test_result'])
                            test_code.append(test_case['test_case'])
                            find_bugs.append(test_case['find_bug'])
                            test_point.append(test_case['test_point'])
                            test_report.append(test_case['test_report'])
                            if test_case['coverage_report']['result'] == 'Success':
                                line_coverage_rates.append(test_case['coverage_report']['output']['line_coverage'])
                                branch_coverage_rates.append(test_case['coverage_report']['output']['branch_coverage'])
                                total_lines.update(test_case['coverage_report']['output']['covered_lines'])
                                total_lines.update(test_case['coverage_report']['output']['missed_lines'])
                                covered_lines.update(test_case['coverage_report']['output']['covered_lines'])
                            else:
                                line_coverage_rates.append(0)
                                branch_coverage_rates.append(0)
                            if test_case['mutation_report']['result'] == 'Success':
                                mutation_scores.append(test_case['mutation_report']['output']['mutation_score'])
                                mutants_list.append(test_case['mutation_report']['output']['filtered_mutations'])
                                for item_mutation in test_case['mutation_report']['output']['filtered_mutations']:
                                    key = str(item_mutation['Line']) + '_' + str(item_mutation['Mutator'])
                                    if key in mutation_info.keys() and mutation_info[key] == 'KILLED':
                                        continue
                                    else:
                                        mutation_info[key] = item_mutation['Status']
                            else:
                                mutation_scores.append(0)

                        logger.info(f"test result: {LoggerManager.format_dict(result)}\n")

                        cal_coverage = len(covered_lines) / len(total_lines) if len(total_lines) != 0 else 0
                        final_coverage = cal_coverage if cal_coverage > max(line_coverage_rates) / 100 else max(
                            line_coverage_rates) / 100
                        logger.info(f"final coverage: {final_coverage}\n")
                        final_mutation_score = sum([1 for v in mutation_info.values() if v == 'KILLED']) / len(
                            mutation_info) if len(mutation_info) != 0 else 0
                        logger.info(f"final mutation score: {final_mutation_score}\n")

                        add_record(log_dir, str(index) + '_' + method.name, result, find_bugs, final_coverage,
                                   final_mutation_score,
                                   line_coverage_rates, branch_coverage_rates, mutation_scores, total_lines,
                                   covered_lines, mutation_info)
                        add_test_case(log_dir, str(index) + '_' + method.name, find_bugs, test_point, test_code,
                                      method.signature, mutants_list)

                    logger_manager.remove_file_handler(local_handler)
                break

            except Exception as e:
                retries += 1
                logger.exception(
                    f"===================== ERROR OCCUR (Retry {retries}/{max_retries}) =====================\n{e}")
                if retries == max_retries:
                    logger.error(f"Max retries reached for method {method}. Skipping...\n")
                    add_record(log_dir, str(index) + '_' + method.name, ["FAIL"], [], 0, 0, [], [], [])


if __name__ == '__main__':
    run_generation()